Greptile Summary

This PR deprecates GSPMD sharding propagation in the JAX backend of TransformerEngine in preparation for its full removal in June 2026. It removes all GSPMD-specific test variants across distributed test suites and example encoders, strips the use_shardy toggle from test helpers and CLI parsers, and adds a version-guarded registration path in BasePrimitive so that:

• On JAX ≤ 0.9.1, infer_sharding_from_operands is still registered (for backwards compatibility) with a one-shot DeprecationWarning.
• On JAX > 0.9.1, the infer_sharding_from_operands kwarg is omitted entirely from def_partition(), since the API was removed from JAX.
• New primitives that never implemented infer_sharding_from_operands inherit a base-class version that raises NotImplementedError with a clear message if GSPMD is attempted.

Key changes:

• transformer_engine/jax/cpp_extensions/base.py: Adds _JAX_GSPMD_SUPPORTED version flag, converts infer_sharding_from_operands from an @abstractmethod to a default @classmethod that raises NotImplementedError, adds _warn_gspmd_deprecation_once() helper, and conditionally passes GSPMD kwargs in register_primitive().
• All distributed test files (test_distributed_fused_attn.py, test_distributed_layernorm.py, test_distributed_layernorm_mlp.py, test_distributed_permutation.py, test_distributed_softmax.py): Remove use_shardy parametrize, jax.config.update("jax_use_shardy_partitioner", ...) calls, and all *_shardy / *_gspmd test methods.
• Encoder examples: Remove --enable-shardy CLI argument, enable_shardy exec helper parameter, and all corresponding *_shardy test cases from both Python files and the shell test runner.

Notable concern: The deprecation warning in _warn_gspmd_deprecation_once() is triggered at primitive registration time (i.e., at import time), not when GSPMD sharding propagation is actually invoked. This means users on JAX ≤ 0.9.1 who have already switched to Shardy will still receive the "Use it at your own risk" DeprecationWarning on every process startup, even though they are not using GSPMD.

Confidence Score: 4/5

• Safe to merge; the GSPMD removal is clean and well-scoped, with one design concern about when the deprecation warning fires.
• All test file cleanups are consistent and complete — no orphaned use_shardy references or dangling jax.config.update calls remain. The version guard in base.py is correctly implemented with PkgVersion and the fallback NotImplementedError for primitives without GSPMD rules is appropriate. Score is 4 rather than 5 because the deprecation warning fires at import time rather than at GSPMD invocation time, which will produce false-positive warnings for users on JAX ≤ 0.9.1 who have already adopted Shardy.
• transformer_engine/jax/cpp_extensions/base.py — specifically the deprecation warning trigger timing in register_primitive().

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[register_primitive called at import time] --> B{JAX version\n≤ 0.9.1?}
    B -- Yes\n_JAX_GSPMD_SUPPORTED=True --> C{cls defines\ninfer_sharding_from_operands\nin own __dict__?}
    C -- Yes --> D[_warn_gspmd_deprecation_once\nDeprecationWarning issued\nat import time]
    D --> E[gspmd_kwargs = infer_sharding_from_operands\nfrom subclass]
    C -- No --> F[gspmd_kwargs = infer_sharding_from_operands\nfrom BasePrimitive base — raises NotImplementedError]
    E --> G[def_partition with partition +\nshardy_sharding_rule + infer_sharding_from_operands]
    F --> G
    B -- No\n_JAX_GSPMD_SUPPORTED=False --> H[gspmd_kwargs = empty dict]
    H --> I[def_partition with only\npartition + shardy_sharding_rule]
    G --> J[Primitive registered]
    I --> J

    style D fill:#ffcc00,color:#000
    style F fill:#ff9999,color:#000

_{Last reviewed commit: 69a71c6}

/te-ci JAX L1

/te-ci JAX L1

Per offline discussion, we decided that GSPMD sharding propagation rules will be kept until June 2026.

This PR only removes the GSPMD tests.
GSPMD rules with the existing primitives will work with older JAX versions (until 0.9.1) with a printed warning.
For the incoming primitives that do not have the GSPMD rules, an Error will be raised only when the primitive is used with GSPMD.

/te-ci JAX L1